# core/async_tools.py
import asyncio
import json
import logging
from fastapi import WebSocket
from typing import Dict, Any

logger = logging.getLogger(__name__)

# =========================
# 🔹 异步任务池管理器
# =========================
class AsyncTaskManager:
    def __init__(self):
        self.tasks = set()

    def create(self, coro):
        """统一管理后台任务，防止内存泄漏"""
        task = asyncio.create_task(coro)
        self.tasks.add(task)
        task.add_done_callback(self.tasks.discard)
        return task


task_manager = AsyncTaskManager()

# =========================
# 🔹 安全发送封装
# =========================
async def safe_send_json(websocket: WebSocket, data: Dict[str, Any]):
    """安全推送 JSON 消息，防止 WebSocket 断开导致异常"""
    try:
        await websocket.send_text(json.dumps(data, ensure_ascii=False))
    except Exception as e:
        logger.warning(f"⚠️ 前端连接中断或发送失败: {e}")


# =========================
# 🔹 工具调用 + 异步总结主逻辑
# =========================
async def async_tool_and_summary(tool_name, tool_args, websocket: WebSocket, ai_core):
    async def run_tool():
        try:
            # 限制工具调用的并发数量
            tool_semaphore = asyncio.Semaphore(3)  # 最多同时执行3个工具调用
            
            async with tool_semaphore:
                # ✅ 立即给用户友好的自然语言反馈，而不是写死的提示
                try:
                    # 让AI生成自然语言的友好回应
                    friendly_prompt = f"用户需要查询{tool_name}工具，请用一句自然友好的中文回应用户，比如'好的，让我帮您查一下'或'请稍等，我马上为您查询'等，不要提及具体工具名称，不要开始执行工具调用"
                    friendly_response = await ai_core.process_query(
                        friendly_prompt,
                        stream=False,
                        enable_tools=False
                    )
                except Exception as e:
                    logger.error(f"生成友好响应时出错: {e}")
                    friendly_response = None
                
                # 如果AI无法生成友好响应，使用默认响应
                if not friendly_response or not friendly_response.strip():
                    friendly_response = "好的，请稍等，我马上为您查询..."
                    logger.info(f"使用默认友好响应: {friendly_response}")
                
                # 发送友好响应给客户端
                if friendly_response:
                    await safe_send_json(websocket, {
                        "type": "friendly_response",
                        "content": friendly_response
                    })
                    logger.info(f"发送友好响应: {friendly_response}")

                # 获取工具描述，生成更友好的提示
                tool_description = ""
                try:
                    # 从ai_core的fastmcp_handler中获取工具描述
                    tool_description = ai_core.fastmcp_handler.get_tool_description(tool_name)
                except Exception as e:
                    logger.error(f"获取工具描述时出错: {e}")
                
                # 根据工具描述生成友好提示
                if tool_description:
                    # 提取工具描述的关键信息，生成自然语言提示
                    if "获取当前日期" in tool_description:
                        friendly_message = "我准备获取当前时间"
                    elif "查询" in tool_description and "车站" in tool_description:
                        friendly_message = "我准备查询车站信息"
                    elif "查询" in tool_description and "余票" in tool_description:
                        friendly_message = "我准备查询车票余票信息"
                    elif "搜索" in tool_description:
                        friendly_message = "我准备搜索相关信息"
                    else:
                        # 默认提示，基于工具描述
                        friendly_message = f"我准备{tool_description[:10]}..."
                else:
                    friendly_message = f"正在查询 {tool_name} ..."
                
                # 通知前端：工具调用开始
                await safe_send_json(websocket, {
                    "type": "tool_start",
                    "tool_name": tool_name,
                    "message": friendly_message
                })

                # 对于新闻查询，优化参数以减少返回数据量
                if tool_name in ["webSearchQuark", "webSearchStd"] and "search_query" in tool_args:
                    # 限制新闻查询返回的结果数量
                    if "count" not in tool_args:
                        tool_args["count"] = 3  # 只返回前3条结果
                    # 设置更严格的时效性过滤
                    if "search_recency_filter" not in tool_args:
                        tool_args["search_recency_filter"] = "oneDay"  # 只查询一天内的新闻
                
                # 对于天气查询，优化参数以提高响应速度
                if tool_name == "webSearchSogou" and "search_query" in tool_args:
                    # 限制天气查询返回的结果数量
                    if "count" not in tool_args:
                        tool_args["count"] = 5  # 只返回前5条结果
                    # 设置更严格的时效性过滤
                    if "search_recency_filter" not in tool_args:
                        tool_args["search_recency_filter"] = "oneDay"  # 只查询一天内的天气信息
                    # 设置较小的内容大小以减少传输时间
                    if "content_size" not in tool_args:
                        tool_args["content_size"] = "high"  # 使用高内容大小以获取足够信息

                # 限时执行工具（防卡死）
                result = await asyncio.wait_for(
                    ai_core.call_mcp_tool(tool_name, tool_args),
                    timeout=15  # 增加超时时间到15秒，因为天气查询可能需要更多时间
                )

                # 推送工具结果
                await safe_send_json(websocket, {
                    "type": "tool_result",
                    "tool_name": tool_name,
                    "result": result
                })
                logger.info(f"🧩 工具执行完成: {tool_name} -> {str(result)[:100]}")

                # 立即并行生成总结，不阻塞主流输出
                task_manager.create(generate_summary(tool_name, result, websocket, ai_core))

        except asyncio.TimeoutError:
            await safe_send_json(websocket, {
                "type": "tool_error",
                "tool_name": tool_name,
                "error": "工具执行超时，请稍后再试"
            })
        except Exception as e:
            await safe_send_json(websocket, {
                "type": "tool_error",
                "tool_name": tool_name,
                "error": str(e)
            })
            logger.error(f"❌ 工具执行异常: {e}")

    # 异步执行
    task_manager.create(run_tool())


# =========================
# 🔹 AI总结逻辑
# =========================
async def generate_summary(tool_name, result, websocket: WebSocket, ai_core):
    try:
        # 优化提示词，使AI更快生成总结
        # 限制结果长度以减少AI处理时间
        result_str = str(result)
        if len(result_str) > 1000:
            result_str = result_str[:1000] + "...(内容已截断)"
        
        prompt = f"请用简洁自然的语言快速总结以下 {tool_name} 结果，适合语音播报，控制在30字以内：\n\n{result_str}"
        
        # 设置更短的超时时间，避免总结过程过长
        try:
            summary = await asyncio.wait_for(
                ai_core.process_query(prompt, stream=False, enable_tools=False),
                timeout=5  # 限制总结过程最多5秒
            )
        except asyncio.TimeoutError:
            summary = "内容总结超时，以下是原始结果摘要："
            # 如果总结超时，提取结果中的关键信息
            if isinstance(result, dict) and "content" in result:
                content = result["content"]
                if isinstance(content, list) and len(content) > 0:
                    first_item = content[0]
                    if isinstance(first_item, dict):
                        title = first_item.get("title", "")
                        if title:
                            summary = f"最新消息：{title}"
                elif isinstance(content, str):
                    # 提取前50个字符作为摘要
                    summary = content[:50] + ("..." if len(content) > 50 else "")

        await safe_send_json(websocket, {
            "type": "ai_summary",
            "tool_name": tool_name,
            "content": summary
        })
        logger.info(f"💬 AI总结完成 ({tool_name}): {summary[:100]}")

        # 发送工具调用完成的通知，确保客户端知道工具调用已经结束
        await safe_send_json(websocket, {
            "type": "tool_complete",
            "tool_name": tool_name,
            "message": f"{tool_name} 工具调用完成"
        })

    except Exception as e:
        await safe_send_json(websocket, {
            "type": "summary_error",
            "tool_name": tool_name,
            "error": str(e)
        })
        logger.error(f"❌ 生成总结异常: {e}")
        
        # 即使总结出错，也要发送工具完成通知
        await safe_send_json(websocket, {
            "type": "tool_complete",
            "tool_name": tool_name,
            "message": f"{tool_name} 工具调用完成（总结生成失败）"
        })